use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};

/// 节点状态
#[derive(Debug, Clone)]
pub struct EndpointHealth {
    pub address: String,                     // 地址
    pub healthy: bool,                       // 健康状态
    pub last_checked: std::time::SystemTime, // 最后检查时间
    pub response_time: Duration,             // 响应时间
}

/// 重试策略
#[derive(Debug, Clone)]
pub enum BalanceStrategy {
    RoundRobin,       // 轮询
    LeastConnections, // 最少连接
}

pub struct LoadBalancer {
    services: Arc<RwLock<HashMap<String, VecDeque<EndpointHealth>>>>,
    strategies: Arc<RwLock<HashMap<String, BalanceStrategy>>>,
}

/// 负载均衡策略
impl LoadBalancer {
    pub fn new() -> Self {
        Self {
            services: Arc::new(RwLock::new(HashMap::new())),
            strategies: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// 添加服务， 如果服务已经存在，则添加新的地址
    pub async fn add_endpoints(&self, service_name: &str, endpoints: Vec<String>) {
        let mut services = self.services.write().await;
        let health_statuses: VecDeque<EndpointHealth> = endpoints
            .into_iter()
            .map(|address| EndpointHealth {
                address,
                healthy: true,
                last_checked: std::time::SystemTime::now(),
                response_time: Duration::from_millis(100),
            })
            .collect();

        services.insert(service_name.to_string(), health_statuses);

        info!("Endpoints added for service: {}", service_name);
    }

    /// 使用其他节点
    pub async fn next_endpoint(&self, service_name: &str) -> Option<String> {
        let mut services = self.services.write().await;

        if let Some(endpoints) = services.get_mut(service_name) {
            if endpoints.is_empty() {
                return None;
            }

            let strategy = self.get_strategy(service_name).await;
            let selected = match strategy {
                BalanceStrategy::RoundRobin => {
                    let endpoint = endpoints.pop_front()?;
                    endpoints.push_back(endpoint.clone());
                    endpoint
                }
                BalanceStrategy::LeastConnections => endpoints
                    .iter()
                    .filter(|ep| ep.healthy)
                    .min_by_key(|ep| ep.response_time)
                    .cloned()?,
            };

            debug!(
                "Selected endpoint: {} for service: {}",
                selected.address, service_name
            );
            Some(selected.address)
        } else {
            None
        }
    }

    /// 获取负载均衡策略
    async fn get_strategy(&self, service_name: &str) -> BalanceStrategy {
        let strategies = self.strategies.read().await;
        strategies
            .get(service_name)
            .cloned()
            .unwrap_or(BalanceStrategy::RoundRobin)
    }

    /// 标记不健康的节点
    pub async fn mark_unhealthy(&self, service_name: &str, endpoint: &str) {
        let mut services = self.services.write().await;

        if let Some(endpoints) = services.get_mut(service_name) {
            for ep in endpoints {
                if ep.address == endpoint {
                    ep.healthy = false;
                    warn!(
                        "Marked endpoint as unhealthy: {} for service: {}",
                        endpoint, service_name
                    );
                }
            }
        }
    }

    /// 标记健康的节点
    pub async fn get_health_status(&self, service_name: &str) -> Vec<EndpointHealth> {
        let services = self.services.read().await;
        services
            .get(service_name)
            .map(|eps| eps.iter().cloned().collect())
            .unwrap_or_default()
    }
}
